package utils;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.baidu.aip.util.Base64Util;
import okhttp3.*;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.net.URLEncoder;
import java.util.List;
import java.util.Map;
import java.util.Scanner;

//对接百度AI平台，图片转文字
public class TestImgToString {
    //设置APPID/AK/SK
    public static final String APP_ID = "29923813";
    public static final String API_KEY = "yuCXaPKn2zPsNnzQPdLs3yRd";
    public static final String SECRET_KEY = "GncL3VBwghFnEFbsT3nAONUGBMqUm41l";
    public static final String ACCESS_TOKEN = "24.136e21e1114f4057174d9555f885c39d.2592000.1677722285.282335-29923813";
    static final OkHttpClient HTTP_CLIENT = new OkHttpClient().newBuilder().build();

    public static void main(String []args) throws IOException{

// 1：根据给定的路径创建一个File对象
        Scanner sc = new Scanner(System.in);
        System.out.println("请输入路径：");
        long start = System.currentTimeMillis();
        String file = sc.nextLine();
        File srcFile = new File(file);
        // 6：调用方法
        getFilePath(srcFile);
        long end = System.currentTimeMillis();
        long time = end-start;
        System.out.println("耗时："+time);
    }
//    循环遍历文件夹，得到图片路径
    public static void getFilePath(File srcFile) {
        // 3：获取给定的File目录下所有的文件或者目录的File数组
        File[] fileArray = srcFile.listFiles();
        // 4：遍历该File数组，得到每一个File对象
        if (fileArray != null) {
            for (File file : fileArray) {
                // 5：判断该File对象是否是目录
                if (file.isDirectory()) {
                    // 5.1：是：递归调用
                    getFilePath(file);
                } else {
                    // 5.2：不是：获取图片绝对路径
                    String imgPath = file.getAbsolutePath();
                    if(!imgPath.contains("txt")){
                        int dot = imgPath.lastIndexOf(".");
                        //得到最后要保存的TXT文件路径
                        String filePath = imgPath.substring(0, dot) + ".txt";
                        boolean flag =  OCRmain(imgPath,filePath);
                        if(flag){
                            System.out.println(filePath+"写入成功");
                        }else{
                            System.out.println(filePath+"写入失败");
                        }
                    }
                }
            }
        }
    }
    //主要步骤
    public static boolean OCRmain(String imgPath,String filePath) {
        boolean paragraph = true;
        //识别
        String res =  accurateBasic(imgPath,paragraph);
        //整理格式，得到文件内容
        String txtContent = getTxtContent(res);
        //写入文件
        return getTxtFile(txtContent,filePath,false);
    }

//    获取access_token
    public static String getAccessToken() throws IOException {
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, "");
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/oauth/2.0/token?client_id="+API_KEY+"&client_secret="+SECRET_KEY+"&grant_type=client_credentials")
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();
        Response response = HTTP_CLIENT.newCall(request).execute();
        String res = response.body().string();
        Map map = JSON.parseObject(res,Map.class);
        return map.get("access_token").toString();
    }

//    通用文字识别（标准版）
    public static String generalBasic(String imgPath,boolean paragraph) {
        // 请求url
        String url = "https://aip.baidubce.com/rest/2.0/ocr/v1/general_basic";
        try {
            // 本地文件路径
            byte[] imgData = FileUtil.readFileByBytes(imgPath);
            String imgStr = Base64Util.encode(imgData);
            String imgParam = URLEncoder.encode(imgStr, "UTF-8");

            String param = "image=" + imgParam + "&paragraph=" + paragraph;

            // 注意这里仅为了简化编码每一次请求都去获取access_token，线上环境access_token有过期时间， 客户端可自行缓存，过期后重新获取。
            String result = HttpUtil.post(url, ACCESS_TOKEN, param);
            System.out.println(result);
            return result;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

//    通用文字识别（高精度版）
    public static String accurateBasic(String imgPath,boolean paragraph) {
        // 请求url
        String url = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate_basic";
        try {
            // 本地文件路径
            byte[] imgData = FileUtil.readFileByBytes(imgPath);
            String imgStr = Base64Util.encode(imgData);
            String imgParam = URLEncoder.encode(imgStr, "UTF-8");
            String param = "image=" + imgParam + "&paragraph=" + paragraph;
            // 注意这里仅为了简化编码每一次请求都去获取access_token，线上环境access_token有过期时间， 客户端可自行缓存，过期后重新获取。
            String result = HttpUtil.post(url, ACCESS_TOKEN, param);
            System.out.println(result);
            return result;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
//  合并段落，得到txt内容
    public static String getTxtContent(String res) {
        StringBuilder txt = new StringBuilder();
        JSONObject resJson = JSON.parseObject(res);
        Object paragraphs_result = resJson.get("paragraphs_result");
        Object words_result = resJson.get("words_result");
        List<Map> paragraphs = ListUtil.castList(paragraphs_result,Map.class);
        List<Map> words = ListUtil.castList(words_result,Map.class);
        List<String> wordList = ListUtil.ListMapToListString(words,"words");
        for(Map idxMap:paragraphs){
            Object idx = idxMap.get("words_result_idx");
            List<Integer> idlist = ListUtil.castList(idx, Integer.class);
            for(int i:idlist){
                txt.append(wordList.get(i));
            }
            txt.append("\n");
        }
//        System.out.println(txt);
        return txt.toString();
    }

    //txt内容写入文件,文件不存在则新建文件
    public static boolean getTxtFile(String txtContent,String filePath,boolean append) {
        boolean flag = false;
        File txtFile = new File(filePath);
        try{
            if(!txtFile.exists()){
                txtFile.createNewFile();//假如文件不存在，新建文件
            }
            FileWriter fw = new FileWriter(filePath,append);
            fw.write(txtContent);
            fw.close();
            flag = true;
        }catch (IOException e){
            e.printStackTrace();
        }
        return flag;
    }


}

